#include "rknn_api.h"
#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <vector>


#include "rk_defines.h"
#include "rk_debug.h"

RK_U8 input_data[1][48][64][1];

static void dump_tensor_attr(rknn_tensor_attr *attr)
{
  char dims[128] = {0};
  for (int i = 0; i < attr->n_dims; ++i)
  {
    int idx = strlen(dims);
    sprintf(&dims[idx], "%d%s", attr->dims[i], (i == attr->n_dims - 1) ? "" : ", ");
  }
  printf("  index=%d, name=%s, n_dims=%d, dims=[%s], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
         "zp=%d, scale=%f\n",
         attr->index, attr->name, attr->n_dims, dims, attr->n_elems, attr->size, get_format_string(attr->fmt),
         get_type_string(attr->type), get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

static inline int64_t getCurrentTimeUs()
{
  struct timeval tv;
  gettimeofday(&tv, NULL);
  return tv.tv_sec * 1000000 + tv.tv_usec;
}

int main(int argc, char *argv[])
{
    printf("==Usage is: model_mfcc modelpath==\n");
    char *model_path = argv[1];
    // char *model_path = "/data/erji-audio-0627.rknn";
    int loop_count = 10;


    rknn_context ctx = 0; // 算法推理上下文
    int ret = rknn_init(&ctx, model_path, 0, 0, NULL); //初始化上下文
    printf("===init model %d===\n", ret);

    rknn_sdk_version sdk_ver; //查询sdk版本号
    ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("rknn_api/rknnrt version: %s, driver version: %s\n", sdk_ver.api_version, sdk_ver.drv_version);

    rknn_input_output_num io_num; //查询输入输出属性
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);


    printf("input tensors:\n");
    //输入tensor的属性
    rknn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        // query info
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret < 0)
        {
            printf("rknn_init error! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&input_attrs[i]);
    }

    printf("output tensors:\n");
    // 输出tensor的属性
    rknn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        // query info
        ret = rknn_query(ctx, RKNN_QUERY_NATIVE_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&output_attrs[i]);
    }

    // set input tensor type
    rknn_tensor_type input_type = RKNN_TENSOR_UINT8;
    rknn_tensor_format input_layout = RKNN_TENSOR_NHWC;
    // Create input tensor memory
    input_attrs[0].type = input_type;
    input_attrs[0].fmt = input_layout;

    rknn_tensor_mem *input_mems[1];
    input_mems[0] = rknn_create_mem(ctx, input_attrs[0].size_with_stride);// size

    // Create output tensor memory
    rknn_tensor_mem *output_mems[1];
    output_mems[0] = rknn_create_mem(ctx, output_attrs[0].size_with_stride);
    //  Set output tensor memory
    output_attrs[0].type = RKNN_TENSOR_FLOAT32;
    if(rknn_set_io_mem(ctx, output_mems[0], &output_attrs[0]) < 0)
    {
        printf("rknn_set_io_mem fail! ret=%d\n", -1);
        return -1;
    }

    RK_U16 numTmp = 0;


    for(int i= 0; i < 48; i++)
    {
        for (int j = 0; j < 64; j++)
        {

            numTmp = (numTmp++ > 255) ? 0 : (numTmp);
            if(numTmp == 256)
            {
                input_data[0][i][j][0] = 0;
            }
            else
            {
                input_data[0][i][j][0] = numTmp;
            }

            // printf("==data is %d===\n", input_data[0][i][j][0]);

            // if(numTmp > 255)
            // {
            //     numTmp = 0;
            //     input_data[0][i][j][0] = 0;
            //     continue;
            // }
            // if(numTmp == 0)
            // {
            //     input_data[0][i][j][0] = 0;
                
            //     numTmp++;
            //     continue;
            // }
            // if(numTmp <= 255)
            // {
            //     input_data[0][i][j][0] = numTmp++;
            // }
            
        }
    }

    // Copy input data to input tensor memory
    int width = input_attrs[0].dims[2];
    int stride = input_attrs[0].w_stride;
    RK_LOGI("====width is %d, stride is %d", width, stride);
    RK_U8 *input_ptr = &input_data[0][0][0][0];
    if(width == stride)
    {
        RK_LOGI("===== == ======\n");
        memcpy(input_mems[0]->virt_addr, input_ptr, width * input_attrs[0].dims[1] * input_attrs[0].dims[3]);
        RK_LOGI("==data length is %d===\n", (width * input_attrs[0].dims[1] * input_attrs[0].dims[3]));

    }else{
        RK_LOGI("===== != ========\n");
        int height = input_attrs[0].dims[1];
        int channel = input_attrs[0].dims[3];
        
        RK_U8 *src_ptr = input_ptr;
        RK_U8 *dst_ptr = (RK_U8 *)input_mems[0]->virt_addr;
        // width-channel elements
        int src_wc_elems = width * channel;
        int dst_wc_elems = stride * channel;
        for (int h = 0; h < height; ++h)
        {
            memcpy(dst_ptr, src_ptr, src_wc_elems);
            src_ptr += src_wc_elems;
            dst_ptr += dst_wc_elems;
        }
    }
    // Set input tensor memory
    if (rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]) < 0)
    {
        RK_LOGE("rknn_set_io_mem fail! ");
        return -1;
    }

    int64_t start_us = getCurrentTimeUs();

    rknn_run(ctx, NULL);
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    printf(" Elapse Time = %.2fms, FPS = %.2f\n", elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);


    return 0;
}